from libs.PipeLine import PipeLine, ScopedTiming
from libs.AIBase import AIBase
from libs.AI2D import Ai2d
import os
import ujson
from media.media import *
import nncase_runtime as nn
import ulab.numpy as np
import image
import gc
import sys
import aidemo
import time

class PersonDetectionApp(AIBase):
    def __init__(self, kmodel_path, model_input_size, labels, anchors,
                 confidence_threshold=0.2, nms_threshold=0.5, nms_option=False,
                 strides=[8,16,32], rgb888p_size=[224,224], display_size=[1920,1080], debug_mode=0):
        super().__init__(kmodel_path, model_input_size, rgb888p_size, debug_mode)
        self.kmodel_path = kmodel_path
        self.model_input_size = model_input_size
        self.labels = labels
        self.anchors = anchors
        self.strides = strides
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.nms_option = nms_option
        self.rgb888p_size = [ALIGN_UP(rgb888p_size[0], 16), rgb888p_size[1]]
        self.display_size = [ALIGN_UP(display_size[0], 16), display_size[1]]
        self.debug_mode = debug_mode
        self.color = [0xFFFF00FF]  # 人体检测框颜色
        self.ai2d = Ai2d(debug_mode)
        self.ai2d.set_ai2d_dtype(nn.ai2d_format.NCHW_FMT, nn.ai2d_format.NCHW_FMT, np.uint8, np.uint8)

    def config_preprocess(self, input_image_size=None):
        with ScopedTiming("PersonDetectionPreprocess", self.debug_mode > 0):
            ai2d_input_size = input_image_size if input_image_size else self.rgb888p_size
            top, bottom, left, right = self.get_padding_param()
            self.ai2d.pad([0, 0, 0, 0, top, bottom, left, right], 0, [0,0,0])
            self.ai2d.resize(nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
            self.ai2d.build([1,3,ai2d_input_size[1],ai2d_input_size[0]],[1,3,self.model_input_size[1],self.model_input_size[0]])

    def postprocess(self, results):
        with ScopedTiming("PersonDetectionPostprocess", self.debug_mode > 0):
            return aicube.anchorbasedet_post_process(
                results[0], results[1], results[2],
                self.model_input_size, self.rgb888p_size,
                self.strides, len(self.labels),
                self.confidence_threshold, self.nms_threshold,
                self.anchors, self.nms_option
            )

    def draw_result(self, pl, dets):
        with ScopedTiming("PersonDetectionDraw", self.debug_mode > 0):
            if dets:
                pl.osd_img.clear()
                for det_box in dets:
                    x1, y1, x2, y2 = det_box[2], det_box[3], det_box[4], det_box[5]
                    w = float(x2 - x1) * self.display_size[0] // self.rgb888p_size[0]
                    h = float(y2 - y1) * self.display_size[1] // self.rgb888p_size[1]
                    x1 = int(x1 * self.display_size[0] // self.rgb888p_size[0])
                    y1 = int(y1 * self.display_size[1] // self.rgb888p_size[1])
                    x2 = int(x2 * self.display_size[0] // self.rgb888p_size[0])
                    y2 = int(y2 * self.display_size[1] // self.rgb888p_size[1])

                    # 过滤掉边缘小目标或异常框
                    if (h < 0.1 * self.display_size[0]):
                        continue
                    if (w < 0.25 * self.display_size[0] and (x1 < 0.03 * self.display_size[0] or x2 > 0.97 * self.display_size[0])):
                        continue
                    if (w < 0.15 * self.display_size[0] and (x1 < 0.01 * self.display_size[0] or x2 > 0.99 * self.display_size[0])):
                        continue

                    # 绘制矩形框与检测标签
                    pl.osd_img.draw_rectangle(x1, y1, int(w), int(h), color=self.color[0], thickness=2)
                    pl.osd_img.draw_string_advanced(x1, y1 - 50, 32, " " + self.labels[det_box[0]] + " " + str(round(det_box[1], 2)), color=self.color[0])
            else:
                pl.osd_img.clear()

    def get_padding_param(self):
        dst_w, dst_h = self.model_input_size
        input_w, input_h = self.rgb888p_size
        ratio = min(dst_w/input_w, dst_h/input_h)
        new_w, new_h = int(input_w*ratio), int(input_h*ratio)
        dw, dh = (dst_w-new_w)//2, (dst_h-new_h)//2
        return dh, dh, dw, dw  # 上下左右填充
